Tanks for pinging us.

@sdvillall and I had a look at the PR and we think that we should be able to add a ROCm-specific environment (e.g., called openfold3-rocm) to the pixi setup. Then users could choose the appropriate environment via

pixi run -e openfold3-rocm run_openfold ...
pixi run -e openfold3-cuda12 run_openfold ...
pixi run -e openfold3-cuda13 run_openfold ...
pixi run -e opnefold3-cpu run_openfold ...

Most of the groundwork should already be there. It's a pity that ROCm is not yet on conda-forge, but we should be able to get everything we need from PyPI (with the ROCm-PyTorch-index).

I will try to cherrypick this PR ontop of the pixi beta branch and see if I can get a working environment. However, I'm not sure I will get around to it this week. Also, I might need help testing as I don't have direct access to AMD accelerators. I will come back to you if that becomes a blocker.

Let me know if that goes in the direction you were thinking @jnwei.

Thank you for the detailed review @jnwei.

In practice, a user of OpenFold may submit queries of many different sequences with different sequence length. In this case, is compilation and caching performed for each unique sequence length? If so, how much extra time to these compilation steps add to the overall worfklow?

Triton JIT-compiles a separate kernel variant for each unique sequence length, since SEQ_LEN is a compile-time constant (tl.constexpr) that gets specialized into the generated GPU code. For a fixed model checkpoint, the head count and embedding dimension are constant, so the main source of new compilations is the sequence length varying across queries.

Triton caches compiled kernels to disk (~/.triton/cache) and reuses them across process restarts. Compilation is therefore a one-time cost per unique sequence length per machine. A cold-cache run takes roughly a few seconds longer than a warm-cache run at the same sequence length. For any repeated workload or subsequent runs, the overhead is zero, and the kernel executes at native GPU speed.

Regarding tl.heuristics: this is unrelated to compilation overhead. It pre-computes three boolean flags at kernel launch time (EVEN_Q, EVEN_KV, EVEN_DIM) and folds them in as compile-time constants to eliminate boundary-check branches when the sequence length divides cleanly into the tile size. It is simply a runtime optimization.

Is there a way to skip the compilation step all together if the user prefers to skip these steps?

Not directly. JIT compilation is intrinsic to how Triton generates GPU-native code and cannot be bypassed. However, since the compiled kernels are cached to disk, the cost is paid once per unique sequence length per machine and never again. Users who need predictable latency from the very first query can pre-warm the cache by running a short dummy forward pass at each expected sequence length before submitting real queries. For any repeated workload, the compilation overhead is zero.

Since the option pip install openfold3[rocm] doesn't add any relevant dependencies, I think we should remove the option from pyproject.toml for now.
One reason we might keep the option is if we can add validation logic to the installation option, to check if pytorch was installed using the extra rocm indices. Perhaps @sdvillal and @Emrys-Merlin will also have thoughts here.

Earlier, I was planning to keep the option as a hook for future PyTorch wheel support, but I decided to remove the empty openfold3[rocm] pip extra in favor of plain pip install openfold3. I added and documented a validate-openfold3-rocm console script to verify the ROCm environment after install.

I agree with @Emrys-Merlin that adding a pixi environment in a follow-up PR would be a cleaner solution. environments/production-amd-linux-64.yml we added in this PR could hopefully serve as a blueprint for the openfold3-rocm pixi environment.

@Emrys-Merlin, happy to help with testing and verification whenever you need access to AMD hardware.

@singagan Thank you for the explanations regarding the compilation and caching of the kernels. I see that the kernel caching will speed up inference long term, so long as the triton cache remains. I wonder if there is a recommended way for working with triton caches for users who may have transient access to compute, for example, if a user runs primarily with AWS instances.

@Emrys-Merlin , your suggestions the specifc pixi environment setups sound great! I think that is a great direction to aim for, we can aim to have a follow up PR that includes RocM.

@singagan I have now had the chance to run the full set of unit tests on an AMD GPU (MI210), and I observe a few failures. I am not sure if you have observed these errors in your testing. I had followed the instructions added to the Installation documentation, and verified that validate-openfold3-rocm checks pass. Perhaps I have made a mistake somewhere.

Deepspeed installation / model configuration issue (openfold3/tests/test_of3_model.py and openfold3/tests/test_kernels.py)

• OpenFold may eventually remove support for deepspeed, which would eventually resolve most of these errors. However, the changes required to remove deepspeed may take some time. Until those changes are merged, I'd like to suggest a few small workarounds:
  • Kernel tests: Add a decorator (similar to @compare_utils.skip_unless_triton_installed() for deepspeed tests that skips the test specific to deepspeed if rocm is detected
  • OF3 model tests: The configuration settings for triton may have to be configured manually, similar to the "reduce model size" settings:

    openfold-3/openfold3/tests/test_of3_model.py

    Lines 49 to 52 in 9879cd8

    	if reduce_model_size:
    	# To avoid memory issues in CI
    	config.architecture.pairformer.no_blocks = 4
    	config.architecture.diffusion_module.diffusion_transformer.no_blocks = 4

- Another option for configuring triton settings is to create a [model_preset](https://github.com/aqlaboratory/openfold-3/blob/main/openfold3/projects/of3_all_atom/config/model_setting_presets.yml) for triton. Basically, you may take the contents of the example `triton.yml` and add them to the `model_settings_presets.yml`, something like this:

Then in test_of3_model.py, the preset can be applied with

project_entry = OF3ProjectEntry()
config = project_entry.get_model_config_with_presets(["eval", "triton"])

And to run inference with the triton kernels, the triton.yml simplifies to

model_update:
  presets:
    - predict
    - triton  # defined in projects.of3_all_atom.model_setting_presets.yml
	- low_mem  # to use low memory settings

LMDB cache writing error ( openfold3/tests/test_lmdb.py)

@jnwei Thank you for the detailed report. I have pushed updates for DeepSpeed-related test on ROCm.

• Added skip_if_rocm() decorator in compare_utils.py and updated skip_unless_ds4s_installed() to also skip on ROCm/HIP, which handles all the test_kernels.py failures automatically
• For test_of3_model.py, followed your suggestion and configured Triton settings manually in run_model, similar to the reduce_model_size pattern. On ROCm, test_shape_small_kernels, test_shape_large_eval, and test_shape_large_bf16_train now use use_triton_triangle_kernels=True instead of use_deepspeed_evo_attention=True, so the tests actually run and validate the Triton path rather than being skipped.

I went with the manual configuration approach rather than the model preset to keep the changes minimal for now.

LMDB cache writing error (test_lmdb.py): Observed this as well. Good to know this is a known issue tied to Python version rather than the AMD build. Happy to defer to the separate PR.

Thank you for adding the configuration changes for test_kernels.py and test_of3_models.py

I see that test_kernels.py on AMD now correctly skips all the tests except the triton tests.

For the test_of3_models.py unit tests, I observe the following errors on all of the tests with the recent additions. My read of the error is that the Linear layer is not configured the fused triton linear modules that are specified in the PR description. Did you encounter this error?

FAILED openfold3/tests/test_of3_model.py::TestOF3Model::test_shape_small_fp32[model=train] - RuntimeError: CUDA error: HIPBLAS_STATUS_INTERNAL_ERROR when calling hipblasLtMatmul with transpose_mat1 1 transpose_ma...

Full stack trace:

openfold3/tests/test_of3_model.py:118: in run_model
    batch, outputs = of3(batch=batch)
                     ^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
openfold3/core/runners/model_runner.py:58: in forward
    return self.model(batch)
           ^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
openfold3/projects/of3_all_atom/model.py:651: in forward
    si_input, si_trunk, zij_trunk = self.run_trunk(
openfold3/projects/of3_all_atom/model.py:200: in run_trunk
    s_input, s_init, z_init = self.input_embedder(
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
openfold3/core/model/feature_embedders/input_embedders.py:128: in forward
    a, _, _, _ = self.atom_attn_enc(
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
openfold3/core/model/layers/sequence_local_atom_attention.py:526: in forward
    ql, cl, plm = checkpoint_section(
openfold3/core/utils/checkpointing.py:152: in checkpoint_section
    return exec(fn, args)
           ^^^^^^^^^^^^^^
openfold3/core/utils/checkpointing.py:146: in exec
    return fn(*a)
           ^^^^^^
openfold3/core/model/layers/sequence_local_atom_attention.py:428: in get_atom_reps
    cl, plm = self.ref_atom_feature_embedder(
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
openfold3/core/model/layers/sequence_local_atom_attention.py:119: in forward
    cl = self.linear_ref_pos(batch["ref_pos"])
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = Linear(in_features=3, out_features=128, bias=False)
input = tensor([[[ 0.8304, -0.3758,  0.8874],
         [-0.8095, -0.6798, -0.2348],
         [-0.2665, -1.6142, -1.8564],
    ...3621, -0.0953,  0.9638],
         [-0.1135, -0.3110, -0.0841],
         [-0.3144,  1.4931,  1.6451]]], device='cuda:0')

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        d = input.dtype
        deepspeed_is_initialized = (
            deepspeed_is_installed and deepspeed.comm.comm.is_initialized()
        )
        if self.precision is not None:
            with torch.amp.autocast("cuda", enabled=False):
                bias = (
                    self.bias.to(dtype=self.precision)
                    if self.bias is not None
                    else None
                )
                return nn.functional.linear(
                    input.to(dtype=self.precision),
                    self.weight.to(dtype=self.precision),
                    bias,
                ).to(dtype=d)

        if d is torch.bfloat16 and not deepspeed_is_initialized:
            with torch.amp.autocast("cuda", enabled=False):
                bias = self.bias.to(dtype=d) if self.bias is not None else None
                return nn.functional.linear(input, self.weight.to(dtype=d), bias)

>       return nn.functional.linear(input, self.weight, self.bias)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E       RuntimeError: CUDA error: HIPBLAS_STATUS_INTERNAL_ERROR when calling hipblasLtMatmul with transpose_mat1 1 transpose_mat2 0 m 128 n 474 k 3 lda 3 ldb 3 ldc 128 abcType 0 computeType 2 scaleType 0

openfold3/core/model/primitives/linear.py:140: RuntimeError

Observed this on specific systems while using hipBLASLt. We switched to rocBLAS, which was already in import_utils.py for inference, but not applied during tests. Added a session-scoped autouse fixture in conftest.py to apply the same setting for the full test suite.